import struct
import numpy as np

class MinstImage(object):
    def __init__(self):
        self.images = []

    def get_images(self):
        return self.images
    
    def parse(self, filePath):
        with open(filePath, "rb") as fImage:
            binImage = fImage.read()
            self.images = []
            offset = 0
            # 4bytes magic number, 4bytes image num, 4bytes row len, 4bytes col len(big endian)
            magicNum, imageNum, rowNum, colNum = struct.unpack_from(">IIII", binImage, offset)
            offset += struct.calcsize(">IIII")
            readImageFmt = ">%dB" %(rowNum*colNum)
            imageSize = struct.calcsize(readImageFmt)
            print("[MinstImage] magicNum %s, imageNum %s, rowNum %s, colNum %s" %(magicNum, imageNum, rowNum, colNum))
            for _ in range(imageNum):
                image = struct.unpack_from(readImageFmt, binImage, offset)
                offset += imageSize
                self.images.append(np.array(image).reshape(rowNum, colNum))

    def get_image_len(self):
        return self.images[0].size

class MinstLabel(object):
    def __init__(self):
        self.labels = []
    
    def get_labels(self):
        return self.labels
    
    def parse(self, filePath):
        with open(filePath, "rb") as fLabel:
            binLabel = fLabel.read()
            self.labels= []
            offset = 0
            # 4bytes magic number, 4bytes image num, 4bytes row len, 4bytes col len(big endian)
            magicNum, labelNum = struct.unpack_from(">II", binLabel, offset)
            offset += struct.calcsize(">II")
            readLabelFmt = ">B" 
            labelSize = struct.calcsize(readLabelFmt)
            print("[MinstLabel] magicNum %s, labelNum %s" %(magicNum, labelNum))
            for _ in range(labelNum):
                label = struct.unpack_from(readLabelFmt, binLabel, offset)[0]
                offset += labelSize
                self.labels.append(np.array([0.9 if i == label else 0.1 for i in range(10)]))

    def get_label_len(self):
        return self.labels[0].size